from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph,START,END
from operator import add
"""
子图--图包含图
1.创建带子图的状态class 属性foo 规约
2.创建5个智能体代理   01、02、l1、r1、r1
3.构建子图
4.构建主图，追加子图
5.编译生成图和运行
"""

class State(TypedDict):
    foo:Annotated[list[int],add]

def agent_01(state):
    return {"foo":["01"]}

def agent_02(state):
    return {"foo":["02"]}

def agent_l1(state):
    return {"foo":["l1"]}

def agent_r1(state):
    return {"foo":["r1"]}

def agent_r2(state):
    return {"foo":["r2"]}

sub_builder = StateGraph(State)
sub_builder.add_node("agent_r1",agent_r1)
sub_builder.add_node("agent_r2",agent_r2)
sub_builder.add_edge(START,"agent_r1")
sub_builder.add_edge("agent_r1","agent_r2")
sub_builder.add_edge("agent_r2",END)
sub_graph = sub_builder.compile()

builder = StateGraph(State)
builder.add_node("agent_01",agent_01)
builder.add_node("agent_l1",agent_l1)
builder.add_node("agent_02",agent_02)
builder.add_node("sub_node",sub_graph) #将节点与子图绑定

builder.add_edge(START,"agent_01")
builder.add_edge("agent_01","agent_l1")
builder.add_edge("agent_l1","agent_02")
builder.add_edge("agent_02",END)

builder.add_edge("agent_01","sub_node")
builder.add_edge("sub_node","agent_02")

builder.add_edge("agent_02",END)

graph = builder.compile()

r = graph.invoke({"foo":[]})
print(r)

graph.get_graph(xray=True).draw_mermaid_png(output_file_path="sub_graph_04.png")

